from typing import AnyStr
from pydantic import BaseModel
import uvicorn
from fastapi import FastAPI,  APIRouter, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import RedirectResponse

from pagecorrection import PageCorrection
from utils import img2base64


app = FastAPI()
cache = {}

router = APIRouter(prefix="/api")


class Picture(BaseModel):
    base64: AnyStr


@router.post("/pic")
async def _(pic: Picture, uuid=Header(None)):
    if not uuid:
        return {"error": True, "msg": "uuid missing", "data": None}
    o = PageCorrection()
    try:
        pages = o.openFromBase64(pic.base64)
    except:
        return {"error": True, "msg": "base64 decode failed", "data": None}
    cache[uuid] = o
    return {"error": False, "msg": "", "data": [img2base64(item) for item in pages]}


@router.post("/cancel")
async def _(uuid=Header(None)):
    if not uuid:
        return {"error": True, "msg": "uuid missing", "data": None}
    if cache.get(uuid):
        cache.pop(uuid)
    return {"error": False, "msg": "", "data": None}


@router.get("/result")
async def _(uuid=Header(None)):
    if not uuid:
        return {"error": True, "msg": "uuid missing", "data": None}
    if not cache.get(uuid):
        return {"error": True, "msg": "no such uuid data", "data": None}
    o = cache[uuid]
    try:
        res = o.process()
        cache.pop(uuid)
    except:
        return {"error": True, "msg": "error occured during processing img", "data": None}
    return {"error": False, "msg": "", "data": [img2base64(item) for item in res]}


app.include_router(router)
app.mount("/app", StaticFiles(directory="./frontend/dist"), name="app")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def _():
    return RedirectResponse("/app/index.html")


if __name__ == "__main__":
    uvicorn.run(app="main:app", host="0.0.0.0",
                port=3000, debug=True, reload=True)
